# -*- coding: utf-8 -*-
"""
提取指定省份的气象观测点经纬度数据

此脚本用于从ERA5气象数据中提取指定省份的观测点经纬度数据，并保存为CSV文件。
使用GIS文件中的省份边界数据来确定每个省份的范围。
"""

import os
import pandas as pd
import numpy as np
import xarray as xr
import argparse
from shapely.geometry import Point
from scripts.gis.province_utils import read_province_shapefile, filter_provinces


# 设置命令行参数
def parse_args():
    parser = argparse.ArgumentParser(description='提取指定省份的气象观测点经纬度数据')
    parser.add_argument('--provinces', type=str, nargs='+', default=None,
                        help='要提取的省份名称列表，如不指定则提取所有省份')
    parser.add_argument('--era5_file', type=str, 
                        default='Daily_ERA5(1960-2024)/ERA5.P.daily.2020.nc',
                        help='ERA5数据文件路径')
    parser.add_argument('--shp_file', type=str, 
                        default='GS(2023)2767审图号/省面.shp',
                        help='省份边界GIS文件路径')
    parser.add_argument('--name_field', type=str, default='NAME',
                        help='省份名称字段')
    parser.add_argument('--output_dir', type=str, default='province_grid_points',
                        help='输出目录')
    return parser.parse_args()

# 读取省份边界数据
def read_province_bounds_from_gis(shp_path, name_field='NAME'):
    """
    从GIS文件中读取省份边界数据
    
    参数:
        shp_path: 省份边界GIS文件路径
        name_field: 省份名称字段，默认为'NAME'
        
    返回:
        包含省份边界数据的GeoDataFrame
    """
    try:
        # 读取shapefile文件
        gdf = read_province_shapefile(shp_path)
        print(f"成功读取省份边界GIS文件，共{len(gdf)}个省份")
        
        # 提取省份名称列表
        province_names = gdf[name_field].tolist()
        print(f"可用省份: {', '.join(province_names)}")
        
        return gdf
    except Exception as e:
        print(f"读取省份边界GIS文件失败: {e}")
        return None

# 读取ERA5数据的经纬度信息
def read_era5_grid(file_path):
    """
    读取ERA5数据的经纬度网格信息
    
    参数:
        file_path: ERA5数据文件路径
        
    返回:
        经度数组和纬度数组
    """
    try:
        ds = xr.open_dataset(file_path)
        print(f"成功读取ERA5数据，维度: {ds.dims}")
        
        # 获取经纬度数组
        lons = ds.lon.values
        lats = ds.lat.values
        
        # 关闭数据集
        ds.close()
        
        print(f"经度范围: {lons.min():.2f} - {lons.max():.2f}, 分辨率: {np.diff(lons)[0]:.2f}")
        print(f"纬度范围: {lats.min():.2f} - {lats.max():.2f}, 分辨率: {np.diff(lats)[0]:.2f}")
        
        return lons, lats
    except Exception as e:
        print(f"读取ERA5数据失败: {e}")
        return None, None

# 提取指定省份的观测点经纬度数据
def extract_province_grid_points(province_gdf, lons, lats, name_field='NAME', provinces=None):
    """
    提取指定省份的观测点经纬度数据
    
    参数:
        province_gdf: 省份边界GeoDataFrame
        lons: 经度数组
        lats: 纬度数组
        name_field: 省份名称字段，默认为'NAME'
        provinces: 要提取的省份名称列表，如为None则提取所有省份
        
    返回:
        包含各省份观测点经纬度数据的字典
    """
    # 如果未指定省份，则使用所有省份
    if provinces is None:
        provinces = province_gdf[name_field].tolist()
    else:
        # 筛选指定的省份
        filtered_gdf = filter_provinces(province_gdf, provinces, name_field)
        provinces = filtered_gdf[name_field].tolist()
    
    if not provinces:
        print("没有有效的省份可以提取")
        return {}
    
    # 创建结果字典
    result = {}
    
    # 对每个省份提取观测点
    for province in provinces:
        # 获取省份几何对象
        province_geom = province_gdf[province_gdf[name_field] == province]['geometry'].iloc[0]
        
        # 获取省份边界范围
        bounds = province_geom.bounds  # 返回(minx, miny, maxx, maxy)
        min_lon, min_lat, max_lon, max_lat = bounds
        
        # 提取在省份范围内的经纬度点
        province_lons = lons[(lons >= min_lon) & (lons <= max_lon)]
        province_lats = lats[(lats >= min_lat) & (lats <= max_lat)]
        
        # 创建经纬度网格点
        lon_grid, lat_grid = np.meshgrid(province_lons, province_lats)
        
        # 将网格点展平为一维数组
        grid_points_flat = np.column_stack((lon_grid.flatten(), lat_grid.flatten()))
        
        # 筛选真正在省份多边形内的点
        in_province_points = []
        for point in grid_points_flat:
            if province_geom.contains(Point(point[0], point[1])):
                in_province_points.append(point)
        
        # 转换为numpy数组
        if in_province_points:
            grid_points = np.array(in_province_points)
        else:
            grid_points = np.empty((0, 2))
        
        # 保存到结果字典
        result[province] = {
            'grid_points': grid_points,
            'lon_count': len(province_lons),
            'lat_count': len(province_lats),
            'total_points': len(grid_points)
        }
        
        print(f"省份: {province}, 经度点数: {len(province_lons)}, 纬度点数: {len(province_lats)}, 有效点数: {len(grid_points)}")
    
    return result

# 保存观测点数据为CSV文件
def save_to_csv(province_data, output_dir):
    """
    保存观测点数据为CSV文件
    
    参数:
        province_data: 包含各省份观测点经纬度数据的字典
        output_dir: 输出目录
    """
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"创建输出目录: {output_dir}")
    
    # 保存汇总信息
    summary_data = []
    for province, data in province_data.items():
        summary_data.append({
            '省份': province,
            '经度点数': data['lon_count'],
            '纬度点数': data['lat_count'],
            '有效点数': data['total_points']
        })
    
    summary_df = pd.DataFrame(summary_data)
    summary_file = os.path.join(output_dir, 'province_grid_points_summary.csv')
    summary_df.to_csv(summary_file, index=False, encoding='utf-8-sig')
    print(f"已保存汇总信息到: {summary_file}")
    
    # 为每个省份保存详细的观测点数据
    for province, data in province_data.items():
        grid_points = data['grid_points']
        if len(grid_points) > 0:
            df = pd.DataFrame(grid_points, columns=['经度', '纬度'])
            
            # 添加省份信息
            df['省份'] = province
            
            # 保存到CSV文件
            file_path = os.path.join(output_dir, f"{province}_grid_points.csv")
            df.to_csv(file_path, index=False, encoding='utf-8-sig')
            print(f"已保存{province}的观测点数据到: {file_path}")
        else:
            print(f"警告: {province}没有有效的观测点数据")
    
    # 保存所有省份的观测点数据到一个文件
    all_points = []
    for province, data in province_data.items():
        grid_points = data['grid_points']
        for point in grid_points:
            all_points.append({
                '省份': province,
                '经度': point[0],
                '纬度': point[1]
            })
    
    if all_points:
        all_df = pd.DataFrame(all_points)
        all_file = os.path.join(output_dir, 'all_provinces_grid_points.csv')
        all_df.to_csv(all_file, index=False, encoding='utf-8-sig')
        print(f"已保存所有省份的观测点数据到: {all_file}")
    else:
        print("警告: 没有有效的观测点数据可以保存")

# 主函数
def main():
    # 解析命令行参数
    args = parse_args()
    
    # 读取省份边界数据
    province_gdf = read_province_bounds_from_gis(args.shp_file, args.name_field)
    if province_gdf is None:
        return
    
    # 读取ERA5数据的经纬度信息
    lons, lats = read_era5_grid(args.era5_file)
    if lons is None or lats is None:
        return
    
    # 提取指定省份的观测点经纬度数据
    province_data = extract_province_grid_points(province_gdf, lons, lats, args.name_field, args.provinces)
    if not province_data:
        return
    
    # 保存观测点数据为CSV文件
    save_to_csv(province_data, args.output_dir)
    
    print("处理完成!")

# 程序入口
if __name__ == '__main__':
    main()